/*
 * Copyright (C) Huawei Technologies Co., Ltd. 2024. All rights reserved.
 * Author: Huawei OS Kernel Lab
 * Create: Fri Apr 19 14:28:18 2024
 */
#include <linux/kobject.h>
#include <linux/module.h>
#include <linux/proc_fs.h>
#include <linux/mm.h>
#include <linux/uaccess.h>
#include <linux/io.h>
#include <linux/platform_device.h>
#include <linux/of.h>
#include <linux/of_address.h>
#include <linux/of_device.h>

#define PO_DEVICE_NAME "post_online_memory"
#define PO_MEMORY_FIELD_NAME "po_memories"
#define MAX_PO_MEMORY  (128)
#define PO_MEMPRY_FILED_NUM (3)
#define MAX_NID_LEN (11)

struct po_memory_st {
	u64	start;
	u64	size;
	u64	nid;
	bool	onlined;
};

static int po_memory_num;
static struct po_memory_st po_memory_table[MAX_PO_MEMORY];

static struct resource isolate_resource = {
	.name	= "isolate mem",
	.start	= 0,
	.end = -1,
	.flags = IORESOURCE_MEM,
};

#define for_each_memory_isolation(r) \
	for(r = isolate_resource.child; r; r = r->sibling)

static void parse_dt_memory_isolations(struct device_node *device_node)
{
	struct device_node *np;
	struct resource r;
	int num_reg = 0;

	np = of_parse_phandle(device_node, "memory-region", 0);

	while (of_address_to_resource(np, num_reg, &r) == 0) {
		__request_region(&isolate_resource, r.start,
			r.end - r.start + 1, "isolate_memory", IORESOURCE_MEM);
		num_reg++;
	}
}

/*
 * format example:
 *
 * 	po_memory {
 * 		compatible = "hisilicon,post-online-memory";
 * 		po_memories =
 * 			<0x00000010 0x80000000 0x1 0x80000000 0 32>; // <start, size, nid> in u64 format
 * 	};
 */
static int po_memory_probe(struct platform_device *pdev)
{
	struct device *dev = &pdev->dev;
	struct device_node *np = dev->of_node;
	int origin_size;
	int i;

	if (!np)
		return -ENODEV;

	if (!of_get_property(np, PO_MEMORY_FIELD_NAME, &origin_size))
		return -EINVAL;

	parse_dt_memory_isolations(np);

	po_memory_num = origin_size / (sizeof(u64) * PO_MEMPRY_FILED_NUM);
	if (po_memory_num > MAX_PO_MEMORY)
		return -E2BIG;

	for (i = 0; i < po_memory_num; i++) {
		int offset = i * PO_MEMPRY_FILED_NUM;
		if (of_property_read_u64_index(np, PO_MEMORY_FIELD_NAME, offset, &po_memory_table[i].start))
			goto err;
		if (of_property_read_u64_index(np, PO_MEMORY_FIELD_NAME, offset + 1, &po_memory_table[i].size))
			goto err;
		if (of_property_read_u64_index(np, PO_MEMORY_FIELD_NAME, offset + 2, &po_memory_table[i].nid))
			goto err;
	}

	for (i = 0; i < po_memory_num; i++) {
		dev_dbg(dev, "[DEBUG] po memory start %lx, size %lx, nid %d, onlined %d.\n",
			po_memory_table[i].start,
			po_memory_table[i].size,
			po_memory_table[i].nid,
			po_memory_table[i].onlined);
	}

	return 0;

err:
	return -EFAULT;
}

static const struct of_device_id po_of_match[] = {
	{ .compatible = "hisilicon,post-online-memory"},
	{ },
};
MODULE_DEVICE_TABLE(of, po_of_match);

static struct platform_driver po_memory_driver = {
	.probe = po_memory_probe,
	.driver = {
		.name = PO_DEVICE_NAME,
		.of_match_table = po_of_match,
	},
};

static struct resource guard = {
	.start = (phys_addr_t)-1, /* MAX */
	.end = (phys_addr_t)-1,
};

static struct resource dummy = {
	.start = 0, /* MIN */
	.end = 0,
};

static struct resource *rsv = &dummy;

static void update_reserve_region(phys_addr_t start)
{
	struct resource *r;

	for_each_memory_isolation(r)
		if (r->end > start)
			break;

	if (!r)
		r = &guard;
	rsv = r;
}


/*
 * start and end are at least 4K aligned
 *
 * this can be:
 *	4K * 0/1 + 8K * 0/1 + 16K * 0/1 + 32K * 0/1 + 64K * 0/1 + 128K * 0/1 +
 *	256K * 0/1 + 512K * 0/1 + 1M * 0/1 + 2M * 0/1 + 4M * 0/1
 */
static void split_and_online_log2_mem(phys_addr_t start, phys_addr_t end)
{
	int start_order, block_order, order;

	start_order = __ffs(start >> PAGE_SHIFT);
	block_order = __fls((end - start) >> PAGE_SHIFT);
	order = min(min(start_order, block_order), MAX_ORDER);

	generic_online_page(phys_to_page(start), order);
	start = start + (PAGE_SIZE << order);
	if (start < end)
		split_and_online_log2_mem(start, end);
}

/*
 * Return:
 *	next start
 */
static phys_addr_t po_memory_split_range(phys_addr_t start, phys_addr_t end)
{
	phys_addr_t next_start;

	if (end < rsv->start) {
		split_and_online_log2_mem(start, end);
		next_start = end;
	} else if (start < rsv->start) {
		split_and_online_log2_mem(start, rsv->start);
		next_start = rsv->start;
	} else if (rsv->start <= start && start < rsv->end) {
		next_start = rsv->end + 1;
		if (end >= rsv->end)
			update_reserve_region(next_start);
	}

	return next_start;
}

static void po_memory_online_page(struct page *page, unsigned int order)
{
	phys_addr_t start, end, next;

	start = page_to_phys(page);
	end = start + (PAGE_SIZE << order);

	while (start < end) {
		/* find next usable address */
		next = po_memory_split_range(start, end);
		start = next;
	}
}

static void online_po_memory(int nid)
{
	int ret, i;
	struct resource *po;
	struct resource *r;

	ret = set_online_page_callback(po_memory_online_page);
	if (ret) {
		pr_err("set callback function failed\n");
		return;
	}

	for (i = 0; i < po_memory_num; i++) {
		if (po_memory_table[i].nid == nid && !po_memory_table[i].onlined) {
			update_reserve_region(po_memory_table[i].start);
			ret = add_memory(nid,
					 po_memory_table[i].start,
					 po_memory_table[i].size,
					 MHP_NONE);
			if (ret) {
				pr_err("po_memory online failed, ret=%d, start %llx, size %llx, nid %lld\n",
				       ret, po_memory_table[i].start, po_memory_table[i].size, po_memory_table[i].nid);
				continue;
			}

			/* show reserved_offline@xxx in /proc/iomem */
			po = lookup_resource(&iomem_resource, po_memory_table[i].start);
			if (!po) {
				pr_err("lookup po memory resource failed\n");
				continue;
			}

			for_each_memory_isolation(r) {
				resource_size_t start, end;

				start = max(PFN_PHYS(PFN_DOWN(r->start)), po->start);
				end = min(PFN_PHYS(PFN_UP(r->end)) - 1, po->end);

				if (start > po->end || end < po->start)
					continue;

				reserve_region_with_split(po, start, end, "reserved");
			}

			po_memory_table[i].onlined = 1;
		}
	}
	restore_online_page_callback(po_memory_online_page);
}

static ssize_t proc_write(struct file *filp, const char __user *buf, size_t len, loff_t *ppos)
{
	int ret;
	char args_buf[MAX_NID_LEN] = {0};
	int nid = 0;

	if (len >= MAX_NID_LEN)
		return -EINVAL;

	ret = copy_from_user(args_buf, buf, len);
	if (ret)
		return -EFAULT;

	ret = kstrtouint(args_buf, 10, &nid);
	if (ret)
		return -EINVAL;

	if (nid > MAX_NUMNODES)
		return -EINVAL;

	online_po_memory(nid);

	return len;
}

static const struct proc_ops proc_ops = {
	.proc_write = proc_write,
};

int __init po_memory_init(void)
{
	int err;

	err = platform_driver_register(&po_memory_driver);
	if (err) {
		pr_err("failed to register po_memory_driver, err=%d.\n", err);
		return err;
	}

	if (!proc_create(PO_DEVICE_NAME, 0200, NULL, &proc_ops)) {
		pr_err("failed to create file /proc/%s.\n", PO_DEVICE_NAME);
		platform_driver_unregister(&po_memory_driver);
		return -EFAULT;
	}

	return 0;
}

void __exit po_memory_exit(void)
{
	remove_proc_entry(PO_DEVICE_NAME, NULL);
	platform_driver_unregister(&po_memory_driver);
}

module_init(po_memory_init);
module_exit(po_memory_exit);
MODULE_LICENSE("GPL");
